import os
import shutil

import fitz
import pymupdf
from PyPDF2 import PdfReader, PdfWriter


def get_file_size(file_path):
    """
    return the size of the file in MB
    """
    return os.path.getsize(file_path) / 1024 / 1024  # in MB


class PyPDF2Helper:

    def get_pdf_page_size(self, pdf_page):
        # save in temp directory
        temp_dir = "temp"
        os.makedirs(temp_dir, exist_ok=True)
        temp_path = os.path.join(temp_dir, "temp.pdf")
        with open(temp_path, "wb") as temp_file:
            writer = PdfWriter()
            writer.add_page(pdf_page)
            writer.write(temp_file)
        size = get_file_size(temp_path)
        return size

    def create_pdf_from_pages(self, pages, output_path):
        """
        Create a new PDF file from a list of PyPDF2 PageObject instances.
        """
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        output_pdf = PdfWriter()
        for page in pages:
            output_pdf.add_page(page)
        with open(output_path, "wb") as output_file:
            output_pdf.write(output_file)

    def split_pdf(self, input_path, output_directory, max_size):
        """
        Given a pdf file, split it into multiple parts if the size exceeds the max_size.\n
        save the parts in the output_directory\n
        the parts are named as <input_file_name>_part1.pdf, <input_file_name>_part2.pdf, ...
        """
        if not input_path.endswith(".pdf"):
            return
        if get_file_size(input_path) <= max_size:
            shutil.copy(input_path, output_directory)
            return
        # Create the output directory if it doesn't exist
        os.makedirs(output_directory, exist_ok=True)

        # Open the input PDF file
        with open(input_path, "rb") as file:
            pdf = PdfReader(file, strict=False)

            # Calculate the total number of pages in the PDF
            total_pages = len(pdf.pages)
            size_per_page = get_file_size(input_path) / total_pages

            # Initialize variables
            part_num = 1
            start_page = 0
            while start_page < total_pages:
                part_pages = int(max_size // size_per_page)
                end_page = min(start_page + part_pages, total_pages)
                current_size = max_size + 1  # Ensure at least one page is added to the output PDF
                while current_size > max_size:
                    tmp_path = os.path.join("temp", f"tmp.pdf")
                    self.create_pdf_from_pages(pdf.pages[start_page:end_page], tmp_path)
                    current_size = get_file_size(tmp_path)
                    if current_size > max_size:
                        end_page = max(start_page + 1, end_page - 10)
                start_page = end_page

                # pdf size meets the requirement
                file_name = os.path.splitext(os.path.basename(input_path))[0]
                shutil.move(tmp_path, os.path.join(output_directory, f"{file_name}_part{part_num}.pdf"))
                part_num += 1


class PymupdfHelper:

    def split_pdf(self, input_path, output_directory, max_size):
        """
        Given a pdf file, split it into multiple parts if the size exceeds the max_size.\n
        save the parts in the output_directory\n
        the parts are named as <input_file_name>_part1.pdf, <input_file_name>_part2.pdf, ...
        """
        if not input_path.endswith(".pdf"):
            return
        if get_file_size(input_path) <= max_size:
            shutil.copy(input_path, output_directory)
            return
        # Create the output directory if it doesn't exist
        os.makedirs(output_directory, exist_ok=True)

        # Open the input PDF file
        pdf = pymupdf.open(input_path)

        # Calculate the total number of pages in the PDF
        total_pages = len(pdf)
        size_per_page = get_file_size(input_path) / total_pages

        # Initialize variables
        part_num = 1
        start_page = 0
        while start_page < total_pages:
            part_pages = int(max_size // size_per_page)
            end_page = min(start_page + part_pages, total_pages)
            current_size = max_size + 1  # Ensure at least one page is added to the output PDF
            delta_page = 10
            while current_size > max_size:
                tmp_path = os.path.join("temp", f"tmp.pdf")
                pages = list(range(start_page, end_page))
                pdf.select(pages)
                pdf.save(tmp_path, 1)
                pdf.close()
                current_size = get_file_size(tmp_path)
                if current_size > max_size:
                    end_page = max(start_page + 1, end_page - delta_page)
                    delta_page = delta_page * 2
                    if end_page == start_page + 1:
                        return
                pdf = pymupdf.open(input_path)
            start_page = end_page

            # pdf size meets the requirement
            file_name = os.path.splitext(os.path.basename(input_path))[0]
            shutil.move(tmp_path, os.path.join(output_directory, f"{file_name}_part{part_num}.pdf"))
            part_num += 1

    def get_page_count(self, input_path):
        pdf = fitz.open(input_path)
        return len(pdf)

def cut_margins(input_path, output_path, top_margin, bottom_margin):
    # Open the PDF file
    pdf = fitz.open(input_path)

    # Iterate over each page
    for page in pdf:
        # Get the original page dimensions
        original_rect = page.rect

        # Calculate the new rectangle dimensions
        new_rect = fitz.Rect(
            original_rect.x0,
            original_rect.y0 + top_margin,
            original_rect.x1,
            original_rect.y1 - bottom_margin,
        )

        # Crop the page using the new rectangle
        page.set_cropbox(new_rect)

    # Save the modified PDF to the output path
    pdf.save(output_path)
    pdf.close()


if __name__ == "__main__":
    # Example usage
    input_file = (
        r"E:\Common\Desktop\LLM\训练数据\raw\算据\考试\文物保护工程从业资格考试\TB_古文化遗址古墓葬_2024.pdf"
    )
    output_directory = "split_parts"
    max_size_mb = 98

    pdf_helper = PymupdfHelper()
    # pdf_helper = PyPDF2Helper()
    pdf_helper.split_pdf(input_file, output_directory, max_size_mb)
    print("PDF split completed.")
